import numpy as np
import numpy.fft as fft
import matplotlib.pyplot as plt
import os


#一步Lax-wendroff格式
def forward_u_LW1(u,c):
    return u[1:-1]-0.5*c*(u[2:]-u[:-2])+0.5*c**2*(u[2:]-2*u[1:-1]+u[:-2])

#精确解
def exact_u(x,t,a):
    return np.sin(2*np.pi*(x-a*t))

#绘图函数
def diag_u_x(u,out_iter,c):
    
    subfolder='u_x_lw1_c={:.2f}_total_t={:.2f}'.format(c,total_t) #子文件夹名
    full_path_dir = os.path.join(os.getcwd(), subfolder) # 拼接当前工作目录和子文件夹路径 
    # 检查子文件夹是否存在  
    if not os.path.exists(full_path_dir):  
        # 如果不存在，则创建子文件夹  
        os.makedirs(full_path_dir)

    t_diag=out_iter*dt
    plt.figure(figsize=(12, 8), dpi=100)  
    plt.plot(x,u)
    plt.xlabel('x')  
    plt.ylabel('u') 
    plt.title('t={:.2f},c={:.2f}'.format(t_diag,c))  
    plt.grid(True)
    # 定义文件名，包含idiag的值  
    filename_base = 'u_x_plot_step_'  
    file_extension = '.png'  
    filename = f"{filename_base}{out_iter}{file_extension}"  
  
    # 拼接完整的文件路径  
    full_path = os.path.join(os.getcwd(), subfolder, filename)  
  
    # 保存散点图到指定路径  
    plt.savefig(full_path)  
    plt.close()

def diag_u_spectrum(u, out_iter, c):
    subfolder = 'u_spectrum_lw1_c={:.2f}'.format(c)  # 子文件夹名
    full_path_dir = os.path.join(os.getcwd(), subfolder)  # 拼接当前工作目录和子文件夹路径

    # 检查子文件夹是否存在，若不存在则创建
    if not os.path.exists(full_path_dir):
        os.makedirs(full_path_dir)

    t_diag = out_iter * dt
    u_spectrum = fft.fft(u)  # 计算u的傅里叶变换（频谱）

    # 取实部（对于实数序列，虚部为复共轭），并只保留非负频率部分
    real_spectrum = np.abs(u_spectrum)[:int(Nx / 2 + 1)]

    freqs = fft.fftfreq(Nx, dx)  # 计算对应的频率
    non_negative_freqs = freqs[:int(Nx / 2 + 1)]  # 只保留非负频率

    plt.figure(figsize=(12, 8), dpi=100)
    plt.plot(non_negative_freqs, real_spectrum)
    plt.xlabel('Frequency')
    plt.ylabel('Amplitude')
    plt.title('t={:.2f}, c={:.2f}'.format(t_diag, c))
    plt.grid(True)
    # 设置频轴范围（根据需要调整）
    min_freq = 0  # 最小频率
    max_freq = 10  # 最大频率

    plt.xlim(min_freq, max_freq)

    filename_base = 'u_spectrum_plot_step_'
    file_extension = '.png'
    filename = f"{filename_base}{out_iter}{file_extension}"

    full_path = os.path.join(os.getcwd(), subfolder, filename)

    plt.savefig(full_path)
    plt.close()

#设置参数
dx=0.02
#Nx = 101
Lx =3
Nx=int(Lx/dx+1.01)
x = np.linspace(0, Lx, Nx)  #linspace生成的点列是包含头和尾的。
dx = x[-1] - x[-2]  #网格长
#dt = 0.01  #时间步长
ndiag=10
a=1
#c=a*dt/dx
c=0.99
dt=c*dx/a
total_t=10
ntime=int(total_t/dt+1.01)
#ntime=2000  #时间步数
print('dx={:},dt={:},a={:},c={:}'.format(dx,dt,a,c))

#初始条件设置
u = np.zeros_like(x)
u = np.sin(2 * np.pi * x)
u[0]=u[-2]
u[-1]=u[1]
#进入主循环
# end_signal = False  #终止判断条件。可以再加一个判断条件。
out_iter = 0    
flag_plot=False
flag_spectrum=True
while  out_iter<ntime:
    if out_iter % ndiag==0 :
        if flag_plot:
            diag_u_x(u,out_iter,c)
        if flag_spectrum:
            diag_u_spectrum(u,out_iter,c)
    out_iter+=1
    u[1:-1]=forward_u_LW1(u,c)
    u[0]=u[-2]
    u[-1]=u[1]

flag_compare_with_exact=True
if flag_compare_with_exact:
    t_array=np.ones_like(x)*out_iter*dt
    u_e=exact_u(x,t_array,a)
    u_h=np.copy(u)

    myerror_sum=np.linalg.norm((u_e-u_h)[:-1],2)*dx
    print(myerror_sum)